"""
Author: kaihonggg
Date: 2025-09-15
Description: 网络设备配置备份
"""

from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
import os
import time
import logging

import pandas as pd
import paramiko

# 指定设备输出编码：utf-8、gbk
OUTPUT_CODE = 'utf-8'

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 定义不同厂商的配置命令
MANUFACTURER_COMMANDS = {
    'huawei': 'screen-length 0 temporary\n display current-configuration\n',
    'h3c': 'screen-length disable\n display current-configuration\n',
    'ruijie': 'terminal length 0\n show running-config\n',
    'cisco': 'terminal length 0\n show running-config\n'
}

def get_device_config(ip, username, password, port=22, manufacturer='h3c', enable_password=None):
    """通过SSH获取设备配置（支持自定义端口和厂商配置）"""
    try:
        # 创建SSH客户端
        ssh_client = paramiko.SSHClient()
        ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        
        # 设置支持的密钥交换算法，优先使用交换机支持的旧算法
        paramiko.Transport._preferred_kex = [
            'diffie-hellman-group-exchange-sha1',
            'diffie-hellman-group14-sha1',
            'diffie-hellman-group1-sha1',
            'diffie-hellman-group-exchange-sha256',
            'ecdh-sha2-nistp256',
            'ecdh-sha2-nistp384',
            'ecdh-sha2-nistp521'
        ]
        
        # 连接设备 - 禁用密钥认证，强制使用密码认证
        ssh_client.connect(
            ip, 
            username=username, 
            password=password, 
            port=port, 
            timeout=10,
            look_for_keys=False,
            allow_agent=False
        )
        remote_conn = ssh_client.invoke_shell()

        # 等待登录
        time.sleep(2)

        # 发送厂商命令
        command = MANUFACTURER_COMMANDS[manufacturer.lower()]  # 确保小写以匹配字典键
        
        # 只有锐捷和思科才需要enable密码
        if manufacturer.lower() in ['ruijie', 'cisco'] and enable_password:
            remote_conn.send('enable\n')
            time.sleep(1)
            remote_conn.send(enable_password + '\n')
            time.sleep(1)

        remote_conn.send(command)
        time.sleep(5)

        # 获取输出
        output = remote_conn.recv(65535).decode(OUTPUT_CODE)

        # 关闭连接
        ssh_client.close()
        return output

    except Exception as e:
        logging.error(f"连接设备 {ip}:{port} 失败: {str(e)}")
        return None

def process_device(ip, username, password, port, save_dir, manufacturer='h3c', enable_password=None):
    """处理单个设备（支持端口和厂商配置）"""
    # 确保manufacturer为小写
    manufacturer = manufacturer.lower()
    logging.info(f"正在获取 {manufacturer} 设备 {ip}:{port} 的配置...")
    config = get_device_config(ip, username, password, port, manufacturer, enable_password)
    if config:
        save_config(ip, config, save_dir)
    else:
        logging.error(f"获取交换机 {ip}:{port} 的配置失败")

def save_config(ip, config, save_dir):
    """保存配置到本地"""
    current_time = datetime.now().strftime('%Y%m%d_%H%M%S')
    filename = f'{save_dir}\\{ip}_{current_time}.log'
    try:
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(config)
        logging.info(f"成功保存配置到 {filename}")
    except Exception as e:
        logging.error(f"保存配置到 {filename} 失败: {str(e)}")

def main():
    """主函数：获取任务名并执行配置备份"""
    print("=" * 50)
    print("网络设备配置备份工具")
    print("=" * 50)
    
    # 获取任务名
    while True:
        task_name = input("请输入任务名称：").strip()
        if task_name:
            invalid_chars = ['\\', '/', ':', '*', '?', '"', '<', '>', '|']
            if any(char in task_name for char in invalid_chars):
                print("任务名不能包含以下字符：\\ / : * ? \" < > |")
                continue
            break
        else:
            print("任务名不能为空，请重新输入！")
    
    # 创建输出目录
    output_dir = 'output'
    save_dir = os.path.join(output_dir, task_name)
    
    # 创建必要的目录（如果不存在）
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"已创建输出目录：{output_dir}")
    
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
        print(f"已创建任务目录：{save_dir}")
    
    print(f"开始执行任务：{task_name}")
    print("-" * 50)
    
    try:
        # 读取CSV文件
        csv_file = 'device_info.csv'
        if not os.path.exists(csv_file):
            print(f"错误：找不到{csv_file}文件，请确保文件存在于当前目录")
            return
            
        df = pd.read_csv(csv_file)
        
        # 建立列名映射
        column_mapping = {}
        for original_col in df.columns:
            clean_col = original_col.strip().lower()
            if 'manufacturer' in clean_col or 'vendor' in clean_col:  # 同时支持manufacturer和vendor列
                column_mapping['manufacturer'] = original_col
            elif clean_col == 'ip':
                column_mapping['ip'] = original_col
            elif clean_col == 'username':
                column_mapping['username'] = original_col
            elif clean_col == 'password':
                column_mapping['password'] = original_col
            elif clean_col == 'port':
                column_mapping['port'] = original_col
            elif 'enable_password' in clean_col:
                column_mapping['enable_password'] = original_col
        
        # 检查必要的列
        required_columns = ['ip', 'username', 'password']
        missing_columns = [col for col in required_columns if col not in column_mapping]
        if missing_columns:
            print(f"错误：{csv_file}缺少必要的列：{', '.join(missing_columns)}")
            return
            
        device_count = len(df)
        print(f"成功读取设备信息，共发现 {device_count} 台设备")
        
        # 使用多线程处理设备
        with ThreadPoolExecutor(max_workers=10) as executor:
            futures = []
            for index, row in df.iterrows():
                # 获取基本参数
                ip = str(row[column_mapping['ip']]).strip()
                username = str(row[column_mapping['username']]).strip()
                password = str(row[column_mapping['password']]).strip()
                
                # 跳过无效数据行
                if not all([ip, username, password]):
                    print(f"警告：跳过第 {index + 1} 行，IP、用户名或密码为空")
                    continue
                
                # 获取可选参数
                port = int(row[column_mapping['port']]) if 'port' in column_mapping and pd.notna(row[column_mapping['port']]) else 22
                manufacturer = str(row[column_mapping['manufacturer']]).strip().lower() if 'manufacturer' in column_mapping and pd.notna(row[column_mapping['manufacturer']]) else 'h3c'
                
                # 验证厂商有效性
                if manufacturer not in MANUFACTURER_COMMANDS:
                    print(f"警告：跳过第 {index + 1} 行，厂商 '{manufacturer}' 无效，有效值为：{list(MANUFACTURER_COMMANDS.keys())}")
                    continue
                
                # 只在锐捷和思科时才读取enable_password
                enable_password = None
                if manufacturer in ['ruijie', 'cisco'] and 'enable_password' in column_mapping and pd.notna(row[column_mapping['enable_password']]):
                    enable_password = str(row[column_mapping['enable_password']]).strip()
                    
                # 如果是锐捷或思科但没有enable密码，给出提示
                if manufacturer in ['ruijie', 'cisco'] and not enable_password:
                    print(f"警告：{manufacturer} 设备 {ip} 未提供enable密码，可能无法获取完整配置")
                    
                # 提交任务到线程池
                future = executor.submit(process_device, ip, username, password, port, save_dir, manufacturer, enable_password)
                futures.append(future)
            
            # 等待所有任务完成
            for future in futures:
                future.result()
                
        print("-" * 50)
        print(f"任务 {task_name} 执行完成！")
        print(f"配置文件已保存到目录：{save_dir}")
        
    except Exception as e:
        logging.error(f"执行过程中出现错误: {str(e)}")
        print(f"执行过程中出现错误: {str(e)}")

if __name__ == '__main__':
    main()